#include <iostream>
#include <fstream>
#include <ctime>
#include "VideoMatting.h"
#include "DataReader.h"

using namespace std;
using namespace cv;

const int WIDTH = 640;
const int HEIGHT = 480;


void printMat2(CvMat* mat, ostream* out = &cout, bool bgr = false) {
	int cn = CV_MAT_CN(mat->type);
	for (int i=0; i<mat->rows; ++i) {
		for (int j=0; j<mat->cols; ++j) {
			CvScalar sca = cvGet2D(mat, i, j);
			if (cn == 3) {
				if (bgr)
					(*out) << "\t(" << sca.val[2] << ", " << sca.val[1] << ", " << sca.val[0] << ")";
				else
					(*out) << "\t(" << sca.val[0] << ", " << sca.val[1] << ", " << sca.val[2] << ")";
			}
			else {
				(*out) << "\t" << sca.val[0];
			}

		}
		(*out) << endl;
	}
}

void printMat2(CvMat* mat, char filename[], bool bgr = false) {
	ofstream fout;
	fout.open(filename);
	printMat2(mat, &fout, bgr);
	fout.close();
}

void removeSmallPart2(CvMat *mask, int range)
{
	int incR[] = {-range, -range, 0, range, range, range, 0, -range};
	int incC[] = {0, range, range, range, 0, -range, -range, -range};

	uchar *maskPtr = mask->data.ptr;

	for (int i = 0; i < mask->rows; ++i)
	{
		for (int j = 0; j < mask->cols; ++j, ++maskPtr)
		{
			uchar &maskVal = *maskPtr;

			int wCnt = 0;
			int bCnt = 0;
			double sumw = 0;
			double sumb = 0;
			for (int k = 0; k < 8; ++k)
			{
				int nr = i + incR[k];
				int nc = j + incC[k];
				if (nr < 0 || nr >= mask->rows || nc < 0 || nc >= mask->cols)
					continue;
				uchar val = CV_MAT_ELEM(*mask, uchar, nr, nc);
				const int valOffset = 1;
				if (val < valOffset)
				{
					++bCnt;
					sumb += val;
				}
				else if (val > 255-valOffset)
				{
					++wCnt;
					sumw += val;
				}
			}
			if (wCnt >= 7 && abs(maskVal - sumw/wCnt) > 20)
			{
				maskVal = sumw / wCnt;
			}
			else if (bCnt >= 7 && abs(maskVal - sumb/bCnt) > 20)
			{
				maskVal = sumb / bCnt;
			}
		}
	}

	// 	cvShowImage("after", mask);
	// 	cvWaitKey(30);
}

/************************************************************************/
/* 1. Remove too far area.
/************************************************************************/
void GenerateMask(CvMat *depthFrame, CvMat *&mask, const CvRect& rect)
{
	mask = cvCreateMat(depthFrame->rows, depthFrame->cols, CV_8U);
	short* p = depthFrame->data.s;
	uchar* maskPtr = mask->data.ptr;

	int **rowSum;
	rowSum = new int*[depthFrame->rows];
	for (int i = 0; i < depthFrame->rows; ++i)
		rowSum[i] = new int[depthFrame->cols];

	for (int i = 0; i < depthFrame->rows; ++i) {
		for (int j = 0; j < depthFrame->cols; ++j, ++p) {
			if (j > 0)
				rowSum[i][j] = rowSum[i][j-1] + (*p);
			else
				rowSum[i][j] = 0;
		}
	}

	long long **allSum;
	allSum = new long long*[depthFrame->rows];
	for (int i = 0; i < depthFrame->rows; ++i)
		allSum[i] = new long long[depthFrame->cols];
	for (int i = 0; i < depthFrame->rows; ++i) {
		for (int j = 0; j < depthFrame->cols; ++j) {
			if (i > 0)
				allSum[i][j] = rowSum[i][j] + allSum[i-1][j];
			else
				allSum[i][j] = rowSum[i][j];
		}
	}

	p = depthFrame->data.s;
	maskPtr = mask->data.ptr;

	//hand
// 	const int LOWER_BOUND = 3500;
// 	const int UPPER_BOUND = 4500;
// 	const int tooFarWidth = 0;
// 	const int REMAIN_UNKNOWN_ROW = 5;
// 	const int REMAIN_UNKNOWN_COL = 5;

// 	const int LOWER_BOUND = 6500;
// 	const int UPPER_BOUND = 7500;
// 	const int tooFarWidth = 1;
// 	const int REMAIN_UNKNOWN_ROW = 5;
// 	const int REMAIN_UNKNOWN_COL = 5;

// 	const int LOWER_BOUND = 2500;
// 	const int UPPER_BOUND = 4500;
// 	const int tooFarWidth = 0;
// 	const int REMAIN_UNKNOWN_ROW = 12;
// 	const int REMAIN_UNKNOWN_COL = 6;

	//people
// 	const int LOWER_BOUND = 1000;
// 	const int UPPER_BOUND = 8000;
// 	const int tooFarWidth = 30;
// 	const int REMAIN_UNKNOWN_ROW = 3;
// 	const int REMAIN_UNKNOWN_COL = 0;

	for (int i=0; i<depthFrame->rows; ++i) {
		for (int j=0; j<depthFrame->cols; ++j, ++p, ++maskPtr) {
			if (i < rect.y || i >= rect.y + rect.height ||
				j < rect.x || j >= rect.x + rect.width) {
					*maskPtr = 0;
					continue;
			}
			if (*p < LOWER_BOUND)
				*maskPtr = 127;
			else if (*p < UPPER_BOUND)
				*maskPtr = 255;
			else
				*maskPtr = 0;

			int boxRow = i + tooFarWidth/2;
			int boxCol = j + tooFarWidth/2;
			if (*p == 0 && i >= tooFarWidth/2 && j >= tooFarWidth/2 &&
					boxRow < depthFrame->rows &&
					boxCol < depthFrame->cols) {
				if (rowSum[boxRow][boxCol] + rowSum[boxRow-tooFarWidth][boxCol-tooFarWidth] ==
						rowSum[boxRow-tooFarWidth][boxCol] + rowSum[boxRow][boxCol-tooFarWidth])
					*maskPtr = 0;
			}
		}
	}

	cvShowImage("original", mask);

	for (int i = 0; i < 3; ++i)
		removeSmallPart2(mask, i+1);
	cvShowImage("remove_first", mask);

	for (int i = 0; i < depthFrame->rows; ++i)
	{
		for (int j=0; j<depthFrame->cols; ++j) {
			uchar& maskVal = CV_MAT_ELEM(*mask, uchar, i, j);
			uchar maskValPrev = 0;

			if (j >= 1)
				maskValPrev = CV_MAT_ELEM(*mask, uchar, i, j-1);

			if (maskValPrev != maskVal)
			{
				int left = REMAIN_UNKNOWN_ROW;
				for (int k = max(j-left, 0); k < min(j+2*REMAIN_UNKNOWN_ROW-left, depthFrame->cols); ++k)
				{
					uchar& val = CV_MAT_ELEM(*mask, uchar, i, k);
//					if (val != 255)
						val = 127;
				}
				j = min(j+2*REMAIN_UNKNOWN_ROW-left, depthFrame->cols);
			}
		}
	}

	cvShowImage("before", mask);

	for (int i = 0; i < depthFrame->cols; ++i)
	{
		for (int j=0; j<depthFrame->rows; ++j) {
			uchar& maskVal = CV_MAT_ELEM(*mask, uchar, j, i);
			uchar maskValPrev = 0;

			if (j >= 1)
				maskValPrev = CV_MAT_ELEM(*mask, uchar, j-1, i);

			if (maskValPrev != maskVal)
			{
				for (int k = max(j-REMAIN_UNKNOWN_COL, 0); k < min(j+REMAIN_UNKNOWN_COL, depthFrame->rows); ++k)
				{
					uchar& val = CV_MAT_ELEM(*mask, uchar, k, i);
//					if (val != 255)
						val = 127;
				}
				j = min(j+REMAIN_UNKNOWN_COL, depthFrame->rows);
			}
		}
	}

	cvShowImage("after", mask);
//	cvWaitKey(0);


	for (int i = 0; i < depthFrame->rows; ++i) {
		delete[] allSum[i];
		delete[] rowSum[i];
	}
	delete[] allSum;
	delete[] rowSum;

}

CvMat * MultiplyMattingMask(CvMat *img, CvMat *mask)
{
	const double eps = 0.000001;
	CvMat *res;
	res = cvCloneMat(img);
	for (int i = 0; i < mask->rows; ++i)
	{
		for (int j = 0; j < mask->cols; ++j)
			if (cvGetReal2D(mask, i, j) < 127 + eps)
				cvSet2D(res, i, j, cvScalar(0));
	}
	return res;
}

CvMat * MultiplyDb(CvMat *img, CvMat *mask)
{
	CvMat *res;
	res = cvCloneMat(img);
	for (int i = 0; i < mask->rows; ++i)
		for (int j = 0; j < mask->cols; ++j)
		{
			CvScalar pixel = cvGet2D(img, i, j);
			double scale = cvGetReal2D(mask, i, j);
			for (int k = 0; k < 4; ++k)
				pixel.val[k] *= scale;
			cvSet2D(res, i, j, pixel);
		}
	return res;
}

CvMat * MultiplyCh(CvMat *img, CvMat *mask)
{
	CvMat *res;
	res = cvCloneMat(img);
	for (int i = 0; i < mask->rows; ++i)
		for (int j = 0; j < mask->cols; ++j)
		{
			CvScalar pixel = cvGet2D(img, i, j);
			double scale = cvGetReal2D(mask, i, j) / 255.;
			for (int k = 0; k < 4; ++k)
				pixel.val[k] *= scale;
			cvSet2D(res, i, j, pixel);
		}
		return res;
}

CvMat *Get8UC3DepthImage(CvMat *depth)
{
	CvMat *res = NULL;
	res = cvCreateMat(depth->rows, depth->cols, CV_8UC3);
	for (int i = 0; i < depth->rows; ++i)
	{
		for (int j = 0; j < depth->cols; ++j)
		{
			double val = cvGetReal2D(depth, i, j)*255/32767;
			cvSet2D(res, i, j, cvScalar(val, val, val));
		}
	}

	return res;
}

CvMat *MergeMats(CvMat *firstMat, CvMat *secondMat)
{
	int cn = CV_MAT_CN(firstMat->type);
	CvMat *firstMatDb = cvCreateMat(firstMat->rows, firstMat->cols, CV_MAKE_TYPE(CV_32F, cn));
	CvMat *secondMatDb = cvCreateMat(secondMat->rows, secondMat->cols, CV_MAKE_TYPE(CV_32F, cn));
	cvConvertScale(firstMat, firstMatDb, 1./32767);
	cvConvertScale(secondMat, secondMatDb, 1./255);

	CvMat *res = NULL;
	res = cvCreateMat(firstMat->rows, firstMat->cols + secondMat->cols, CV_MAKE_TYPE(CV_32F, cn));
	
	for (int i = 0; i < firstMat->rows; ++i)
	{
		for (int j = 0; j < firstMat->cols; ++j)
			cvSet2D(res, i, j, cvGet2D(firstMatDb, i, j));
		for (int j = 0; j < secondMat->cols; ++j)
			cvSet2D(res, i, firstMat->cols + j, cvGet2D(secondMatDb, i, j));
	}

	return res;
}

int main(int argc, char *argv[])
{
	ColorReader colorReader;
	DepthReader depthReader;

	CvMat *depthFrame = cvCreateMat(HEIGHT, WIDTH, CV_16SC1);
	CvMat *colorFrame = cvCreateMat(HEIGHT, WIDTH, CV_8UC3);

	for (int i = 0; i < 1; ++i)
	{
		depthReader.getNextFrame(depthFrame->data.s);
		colorReader.getNextFrame(colorFrame->data.ptr);
	}

	for (int i = 0; i < 10; ++i)
	{
		depthReader.getNextFrame(depthFrame->data.s);
		colorReader.getNextFrame(colorFrame->data.ptr);
		CvRect rect;
		rect.x = 270;
		rect.y = 100;
		rect.width = 460 - rect.x;
		rect.height = 330 - rect.y;
		rect.x = 0;
		rect.y = 0;
		rect.width = depthFrame->cols;
		rect.height = depthFrame->rows;

		cvRectangleR(depthFrame, rect, cvScalar(32767));
		cvRectangleR(colorFrame, rect, cvScalar(255,255,255));

		cvShowImage("depth", depthFrame);
		cvShowImage("color", colorFrame);

		double totalTime = -clock();
		CvMat *mask = NULL;
		GenerateMask(depthFrame, mask, rect);
		cvShowImage("mask", mask);
		cvShowImage("ori_fg", MultiplyCh(colorFrame, mask));

		CvMat *mattingImg = sharedMatting(colorFrame, mask, true);
		cout << "running time: " << (totalTime + clock()) / CLOCKS_PER_SEC << endl;
		cvShowImage("matting", mattingImg);
		cvShowImage("fg", MultiplyDb(colorFrame, mattingImg));

		cvWaitKey(0);
	}

 	cvWaitKey(0);

	cvDestroyAllWindows();

	return 0;
}
